__doc__ =   """ Contains classes to compute distance contacts for
                residues and atoms.
            """

from numpy import array, r_, sqrt
from collections import defaultdict

from AbstractPropertyDict import AbstractAtomPropertyDict, AbstractResiduePropertyDict,\
                                 AbstractPropertyDict
from ZenExceptions import ZenExternalException, ZenException, ZenImplementationException
from Math import coords_to_symmetry, coords_to_crystal
from ResidueProperties import AA_NAMES_STR
from Externals import call_act, id_hetero
from Crunchers import cnt_loop
from Einput import einput

def get_contacts(structure, query_ents, level, search_limit,\
                 contact_mode, symmetry_mode, crystal_mode,\
                 bucket_size, MAXSYM, MAXCNT):
    """ Searches distance contacts between entities.
        This function searches contacts for query entities (query_ents) in the
        asymmetric unit (ASU, structure). The asymmetric unit can be expanded into
        the biological molecule or unit cell (symmetry_mode) or crystal (crystal_mode).

        @type structure: Structure entity
        @param structure: The contents of the asymmetric unit
        @type query_ents: Collection of entities
        @param query_ents: Query entities for which contacts will be calculated
        @type level: SMRCA hierarchy level string
        @param level: Defines the entity level to search for contacts
        @type search_limit: float distance in A^2
        @param search_limit: maximum distance threshold in square Angstrom
        @type contact_mode: string, "diff_cell", "diff_sym", "diff_chain"
        @param contact_mode: Defines the the allowed contacts i.e. requires
                             that contacts are made by entities:
                             diff_cell - from different unit cells
                             diff_sym - with different symmetry operators
                                        (if in the same unit cell)
                             diff_chain - with differen chain ids
                                        (if in the same unit cell and symmetry)
        @type symmetry_mode: None or one of "table", "raw", "bio"
        @param symmetry_mode: Defines which symmetry matrices should be applied
                             to produce either the biological molecule or unit cell.
                              None - No symmetry matrices
                             "raw" - uses the matrices from the PDB file header to
                                     produce the unit cell
                             "table" - uses the matrices from the crystallographic
                                       tables to produce the unit cell
                             "bio" - creates the biological unit's based on the BIOMT
                                     records in the PDB file header
        @type crystal_mode: int
        @param crystal_mode: The number of unit cells in each direction i.e. 1 is makes
                             a total of 27 unit cells: (-1, 0, 1) == 3, 3^3 == 27

    """
    lattice_ents = structure[(0,)]
    query_ents = (query_ents or lattice_ents)

    sh = structure.header
    contact_mode = {'diff_asu'  :0,
                    'diff_sym'  :1,
                    'diff_chain':2 }[contact_mode]

    lents = einput(lattice_ents, level)
    lents_ids = lents._data_children('get_full_id', forgiving =False, function =True)
    lcoords = array(lents._data_children('coords', forgiving =False))

    qents = einput(query_ents, level)
    qents_ids = qents._data_children('get_full_id', forgiving =False, function =True)
    qcoords = array(qents._data_children('coords', forgiving =False))

    if symmetry_mode:
        if symmetry_mode == 'table':
            lcoords = coords_to_symmetry(lcoords, sh['table_fmx'], sh['table_omx'], sh['table_mxs'], symmetry_mode)
        elif symmetry_mode == 'raw':
            lcoords = coords_to_symmetry(lcoords, sh['raw_fmx'], sh['raw_omx'], sh['raw_mxs'], symmetry_mode)
        elif symmetry_mode == 'bio':
            #lcoords = coords_to_symmetry(lcoords, sh['raw_fmx'], sh['raw_omx'], sh['bio_mxs'], symmetry_mode)
            raise ZenImplementationException, "Not implemented symmetry_mode: %s" % symmetry_mode
        else:
            raise ZenException, "Unsupported symmetry_mode: %s" % symmetry_mode
    else:
        lcoords = array([lcoords]) # fake 3D
    if crystal_mode:
        zero_tra = {1:13, 2:62, 3:171}[crystal_mode]
        # 0,0,0 translation is: Thickened cube numbers: a(n)=n*(n^2+(n-1)^2)+(n-1)*2*n*(n-1).
        # 1, 14, 63, 172, 365, 666, 1099, 1688, 2457, 3430, 4631, 6084, 7813 ...
        # python indexes -> -1
        if symmetry_mode == 'table':
            lcoords = coords_to_crystal(lcoords, sh['table_fmx'], sh['table_omx'], crystal_mode)
        elif symmetry_mode == 'raw':
            lcoords = coords_to_crystal(lcoords, sh['raw_fmx'], sh['raw_omx'], crystal_mode)
        else:
            raise ZenException, "Expansion to crystal not possible for this symmetry_mode: %s" % symmetry_mode
    else:
        zero_tra = 0
        lcoords = array([lcoords]) # fake 4D
    shape = lcoords.shape
    lcoords = lcoords.reshape((shape[0] * shape[1] * shape[2], shape[3]))
    box = r_[qcoords.min(axis=0) - search_limit, qcoords.max(axis=0) + search_limit]

    lc = [] # lattice chain
    qc = [] # query chain
    lchains = [i[2] for i in lents_ids]
    qchains = [i[2] for i in qents_ids]
    allchains = set()
    allchains.update(lchains)
    allchains.update(qchains)
    chain2id = dict(zip(allchains, range(len(allchains))))
    for lent_id in lents_ids:
        lc.append(chain2id[lent_id[2]])
    for qent_id in qents_ids:
        qc.append(chain2id[qent_id[2]])
    lc = array(lc)
    qc = array(qc)

    (idxc, n_src, n_asu, n_sym, n_tra, n_dst) = cnt_loop(qcoords, lcoords, qc, lc, shape[1],\
                                        shape[2], zero_tra, contact_mode, search_limit,\
                                        box, bucket_size, MAXSYM, MAXCNT)
    result = defaultdict(dict)
    for contact in xrange(idxc):
        qent_id = qents_ids[n_src[contact]]
        lent_id = lents_ids[n_asu[contact]]
        result[qent_id][lent_id] = (sqrt(n_dst[contact]), n_tra[contact], n_sym[contact])
    return result



class AbstractContacts(AbstractAtomPropertyDict):
    """ Abstract class to map atom-atom contact onto atom entites. An atom's
        xtra dict. will contain data from the get_contacts output.
    """
    def __init__(self, structure, xtra_key, cnt_args):
        # these are some sane defaults
        # get_contacts(structure, query_ents, level, search_limit,\
        # contact_mode, symmetry_mode, crystal_mode,\
        # bucket_size, MAXSYM, MAXCNT):
        cnt_args['query_ents'] = (cnt_args.get('query_ents') or None)
        cnt_args['search_limit'] = (cnt_args.get('search_limit') or 5.0)
        cnt_args['bucket_size'] = (cnt_args.get('bucket_size') or 10)
        cnt_args['level'] = (cnt_args.get('level') or 'A')
        cnt_args['MAXCNT'] = (cnt_args.get('MAXCNT') or 100000)
        cnt_args['MAXSYM'] = (cnt_args.get('MAXSYM') or 200000)

        contacts = get_contacts(structure, **cnt_args)
        for query_id, cnts in contacts.iteritems():
            entity = structure.table[cnt_args['level']][query_id]
            entity.xtra[xtra_key] = cnts
        AbstractAtomPropertyDict.__init__(self, contacts)


class ASUContacts(AbstractContacts):
    def __init__(self, structure, xtra_key =None, **kwargs):
        xtra_key = (xtra_key or 'aCNT')
        cnt_args = {"contact_mode": "diff_chain",
                    "symmetry_mode": None,
                    "crystal_mode": 0}
        cnt_args.update(kwargs)
        AbstractContacts.__init__(self, structure, xtra_key, cnt_args)


class TotalContacts(AbstractContacts):
    def __init__(self, structure, xtra_key =None, **kwargs):
        xtra_key = (xtra_key or 'tCNT')
        cnt_args = {"contact_mode": "diff_chain",
                    "symmetry_mode": "raw",
                    "crystal_mode": 2}
        cnt_args.update(kwargs)
        AbstractContacts.__init__(self, structure, xtra_key, cnt_args)


#class CrystalContacts(AbstractContacts):
    #def __init__(self, structure, xtra_key =None, **kwargs):
        #xtra_key = (xtra_key or 'CNT')
        #sym_args = {'mode':'table'}
        #cry_args = {'n':2}
        #cnt_args = {'query_ents':None, 'lattice_ents':None,  'mode':'diff_sym',
                    #'expand_symmetry':True, 'expand_crystal':True}
        #cnt_args.update(kwargs)
        #AbstractContacts.__init__(self, structure, xtra_key, sym_args, cry_args, cnt_args)


#class BioContacts(AbstractContacts):
    #def __init__(self, structure, xtra_key =None, **kwargs):
        #xtra_key = (xtra_key or 'CNT')
        #sym_args = {'raw':True, 'bio':True}
        #cnt_args = {'query_ents':None, 'lattice_ents':None,  'mode':'diff_chain',
                    #'expand_symmetry':True, 'expand_crystal':False}
        #cnt_args.update(kwargs)
        #AbstractContacts.__init__(self, structure, xtra_key, sym_args, {}, cnt_args)


class AbstractActContacts(AbstractAtomPropertyDict):
    """ Abaract class to run ncont with kwargs settings
        and map the results onto atoms. An atoms's xtra dict
        will contain a dict of other atom id's in which it
        is in contact along with their distances.
    """
    def __init__(self, pdb_file, structure, app, xtra_key, **kwargs):
        # we do not check if the provided pdb_file and models fit
        # neither if this is a X-ray structure.
        contacts = eval(app)(pdb_file, **kwargs)
        contact_dict = {}
        #model.set_table()
        for atom_full_id in contacts.iterkeys(): # left contacting atom
            if structure:
                structure_id = structure.get_id()
                atom_fuller_id = structure_id + atom_full_id
                atom_contacts = dict([(structure_id + id, (dist, symm)) for (id, dist, symm) in contacts[atom_full_id]])
            else:
                atom_fuller_id = atom_full_id
                atom_contacts = dict([(id, (dist, symm)) for (id, dist, symm) in contacts[atom_full_id]])
            try:
                atom = structure.table['A'][atom_fuller_id]
            except KeyError:
                try:
                    hetero_id = id_hetero(atom_fuller_id, -2)
                    atom = structure.table['A'][hetero_id]
                    atom_fuller_id = hetero_id
                except KeyError:
                    raise PDBExternalException, (atom_fuller_id, hetero_id)
            atom.xtra[xtra_key] = atom_contacts
            contact_dict[atom_fuller_id] = atom_contacts
        AbstractAtomPropertyDict.__init__(self, contact_dict)


class ACTAtomContacts(AbstractActContacts):
    """ Class to search for atoms contacting with atoms of other symmetry mates.
    """
    def __init__(self, pdb_file, structure, xtra_key =None, **kwargs):
        symmetry = structure.header['space_group'].replace(' ','')
        kwargs['symm'] = symmetry
        kwargs['contact'] = (kwargs.get('contact') or "inter 0.0 5.0")
        xtra_key = (xtra_key or 'CNT_ACT_X')
        AbstractActContacts.__init__(self, pdb_file, structure, 'call_act', xtra_key, **kwargs)


#EOF